/*
* Copyright 1993-2010 NVIDIA Corporation.  All rights reserved.
*
* NVIDIA Corporation and its licensors retain all intellectual property and 
* proprietary rights in and to this software and related documentation. 
* Any use, reproduction, disclosure, or distribution of this software 
* and related documentation without an express license agreement from
* NVIDIA Corporation is strictly prohibited.
*
* Please refer to the applicable NVIDIA end user license agreement (EULA) 
* associated with this source code for terms and conditions that govern 
* your use of this NVIDIA software.
* 
*/


#include "common/book.h"
#include "common/gpu_anim.h"

#include <cuda.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <cmath>
#include <string.h>
#include <stdlib.h>

#define DIM		512
#define CANVDIM	512
#define CELLDIM	CANVDIM/DIM
#define GAP		0
#define R		2
#define ANIMATE true
#define MAXREP	pow((float)2,31)

#define AGENTNO DIM*DIM/2
#define EMPTY	-1
int		EXITNO =4;

#define TWO_TO_ONE(x,y) (x)+(y)*DIM
#define POS_TO_EXIT_1D(pos, exit) abs(pos % DIM - exit % DIM) + abs(pos / DIM - exit / DIM)
#define POS_TO_EXIT_2D(x, y, exit) (x-exit%DIM) * (x-exit%DIM) + (y-exit/DIM) * (y-exit/DIM)
//#define POS_TO_EXIT_2D(x, y, exit) abs(x - exit % DIM) + abs(y - exit % DIM)

#define timer_start(start, stop) \
	cudaEventCreate(&start);\
	cudaEventCreate(&stop);\
	cudaEventRecord(start, 0);
#define timer_stop(start, stop) \
	float elapsedTime = -1; \
	cudaEventRecord(stop, 0);\
	cudaEventSynchronize(stop);\
	cudaEventElapsedTime(&elapsedTime, start, stop);\
	cudaEventDestroy(start);\
	cudaEventDestroy(stop);\
	printf("%1.3fms\n", elapsedTime);
	//cout<<elapsedTime<<"ms"<<endl;

__global__ void init(int* state1, float *rands);
__global__ void inspection_kernel (uchar4 *ptr, int *env, int *curpos, int *nextpos, float *rands);
void generate_frame( uchar4 *pixels, void*, int* env, int *exitPos, int* curPos, int* nextPos, float *rands );
void init_state(GPUAnimBitmap &bitmap);
void decision_point(int x, int y, int* nextPos, int *exitPos);

int main( void ) {
	unsigned int i = (unsigned int)pow((float)2, (32-2*log((float)DIM)/log((float)2)));
	printf("%d\n", i);
	GPUAnimBitmap  bitmap( DIM*CELLDIM, DIM*CELLDIM, AGENTNO, NULL );
	bitmap.animated = ANIMATE;
	init_state(bitmap);
	bitmap.anim_and_exit(
		(void (*)(uchar4*,void*,int*,const int*,int*,int*,float*))generate_frame, 
		NULL, 
		(void (*)(int,int,int*,int*))decision_point ); 
	return 0;
}

__global__ void wash(uchar4 *ptr){
	int x = threadIdx.x + blockIdx.x*blockDim.x;
	int y = threadIdx.y + blockIdx.y*blockDim.y;
	int offset = x + y * blockDim.x * gridDim.x;
	uchar4 color;
	color.x = 0;
	color.y = 0;
	color.z = 0;
	color.w = 255;
	if(offset<DIM*DIM*CELLDIM*CELLDIM)
		ptr[offset]=color;		
}

__global__ void paint(uchar4 *canvas, int *curPos, int *nextPos, const int *exitPos){
	int idx = threadIdx.x+blockIdx.x*blockDim.x;
	uchar4 green, yellow, red;
	green.x = 0;
	green.y = 255;
	green.z = 0;
	green.w = 255;
	yellow.x = 255;
	yellow.y = 255;
	yellow.z = 0;
	yellow.w = 255;
	red.x = 255;
	red.y = 0;
	red.z = 0;
	red.w = 0;
	if(idx<AGENTNO){
		int cpos = curPos[idx];
		int npos = nextPos[idx];
		if(cpos>=0){
			int cx = (int)cpos%DIM;
			int cy = (int)cpos/DIM;
			for(int i=0; i<CELLDIM-GAP; i++){
				for(int j=0; j<CELLDIM-GAP; j++){
					int offset = (CELLDIM*cx+i) + (CELLDIM*(DIM-1-cy)+j)*DIM*CELLDIM;
					canvas[offset]=green;
				}
			}
			cpos = exitPos[idx];
			cx = (int)cpos%DIM;
			cy = (int)cpos/DIM;
			for(int i=0; i<CELLDIM-GAP; i++){
				for(int j=0; j<CELLDIM-GAP; j++){
					int offset = (CELLDIM*cx+i) + (CELLDIM*(DIM-1-cy)+j)*DIM*CELLDIM;
					canvas[offset]=red;
				}
			}
		}
	}
}

__device__ bool combat(const int *env, const int agId, const int agNextPos, 
	const int *nextPos, const float *rands){
		int rivalNextPos = -1;
		int rivalId = -1;
		int agNextX = agNextPos%DIM;
		int agNextY = agNextPos/DIM;
		if(agNextX-1 >= 0){
			rivalId = env[agNextPos-1];
			if(rivalId >= 0 && rivalId != agId){
				rivalNextPos = nextPos[rivalId];
				if(rivalNextPos == agNextPos && rands[agId] <= rands[rivalId])
					return false;
			} 
		}
		if(agNextX+1 < DIM){
			rivalId = env[agNextPos+1];
			if(rivalId >= 0 && rivalId != agId){
				rivalNextPos = nextPos[rivalId];
				if(rivalNextPos == agNextPos && rands[agId] <= rands[rivalId])
					return false;
			} 
		}
		if(agNextY-1 >= 0){
			rivalId = env[agNextPos-DIM];
			if(rivalId >= 0 && rivalId != agId){
				rivalNextPos = nextPos[rivalId];
				if(rivalNextPos == agNextPos && rands[agId] <= rands[rivalId])
					return false;
			} 
		}
		if(agNextY+1 < DIM){
			rivalId = env[agNextPos+DIM];
			if(rivalId >= 0 && rivalId != agId){//
				rivalNextPos = nextPos[rivalId];
				if(rivalNextPos == agNextPos && rands[agId] <= rands[rivalId])
					return false;
			}
		}
		return true;
}

__device__ int nextstep(const int *env, const int agExitPos, const int agCurPos){
	int agNextPos = agCurPos;
	int agCurX = agCurPos%DIM;
	int agCurY = agCurPos/DIM;
	int agCurDistToExit = POS_TO_EXIT_2D(agCurX, agCurY, agExitPos);
	int agNextDistToExit = agCurDistToExit;
	int t = 0;
	if(agCurY-1 >= 0){
		if(env[agCurPos-DIM] == EMPTY){
			agNextDistToExit = POS_TO_EXIT_2D(agCurX, agCurY-1, agExitPos);
			t = TWO_TO_ONE(agCurX, agCurY-1);
			if(agNextDistToExit<agCurDistToExit){
				agNextPos = t;
				agCurDistToExit = agNextDistToExit;
			}
		}
	}
	if(agCurY+1 < DIM){
		if(env[agCurPos+DIM] == EMPTY){
			agNextDistToExit = POS_TO_EXIT_2D(agCurX, agCurY+1, agExitPos);
			t = TWO_TO_ONE(agCurX, agCurY+1);
			if(agNextDistToExit<agCurDistToExit){
				agNextPos = t;
				agCurDistToExit = agNextDistToExit;
			}
		}
	}
	if(agCurX-1 >= 0){
		if(env[agCurPos-1] == EMPTY){
			agNextDistToExit = POS_TO_EXIT_2D(agCurX-1, agCurY, agExitPos);
			t = TWO_TO_ONE(agCurX-1, agCurY);
			if(agNextDistToExit<agCurDistToExit){
				agNextPos = t;
				agCurDistToExit = agNextDistToExit;
			}
		}
	}
	if(agCurX+1 < DIM){
		if(env[agCurPos+1] == EMPTY){
			agNextDistToExit = POS_TO_EXIT_2D(agCurX+1, agCurY, agExitPos);
			t = TWO_TO_ONE(agCurX+1, agCurY);
			if(agNextDistToExit<agCurDistToExit){
				agNextPos = t;
				agCurDistToExit = agNextDistToExit;
			}
		}
	}
	return agNextPos;
}

extern __shared__ int smem[];
__global__ void inspection_kernel (const int *env, const int *curPos, int *nextPos, const float *rands){
	int t = 0;
	int tid = threadIdx.x;
	int idx = threadIdx.x + blockIdx.x * blockDim.x;
	if(idx<AGENTNO){
		int agCurPos = curPos[idx];
		int agNextPos = nextPos[idx];
		bool settledown = false;
		if (env[agNextPos] == EMPTY)
			settledown = combat(env, idx, agNextPos, nextPos, rands);
		if (settledown != true)
			agNextPos = agCurPos;
		nextPos[idx]=agNextPos;
	}
}

__global__ void nextstep_kernel (const int *env, const int *exitPos, int *curPos, int *nextPos) {
	int idx = threadIdx.x + blockIdx.x * blockDim.x;
	if(idx<AGENTNO){
		int agCurPos = curPos[idx];
		int agExitPos = exitPos[idx];
		int agNextPos = nextstep(env, agExitPos, agCurPos);
		nextPos[idx] = agNextPos;
	}
}

__global__ void register_kernel (int *env, const int *exitPos, int *curPos, const int *nextPos){
	int idx = threadIdx.x+blockIdx.x*blockDim.x;
	if(idx<AGENTNO){
		int t = curPos[idx];
		if (t >= 0 && t < DIM * DIM)
			env[t] = EMPTY;
		t = nextPos[idx];
		curPos[idx] = t;
		env[t] = idx;
		if(t==exitPos[idx])
			env[t]=EMPTY;
	}
}

CUevent start, stop;
int newStateCountDown = 100;
void generate_frame( uchar4 *canvas, void*, int* env, int *exitPos, int* curPos, int* nextPos, float *rands ) {
	//timer_start(start, stop);
	
	dim3 washblock(32,32);
	dim3 washgrid(DIM*CELLDIM/32,DIM*CELLDIM/32);
	int blockSize = 32;
	int gridSize = AGENTNO/blockSize+1;

	register_kernel<<<gridSize, blockSize>>>(env, exitPos, curPos, nextPos);
	nextstep_kernel<<<gridSize, blockSize>>>(env, exitPos, curPos, nextPos);
	inspection_kernel<<<gridSize, blockSize>>>(env, curPos, nextPos, rands);
	wash<<<washgrid, washblock>>>(canvas);
	paint<<<gridSize, blockSize>>>(canvas, curPos, nextPos, exitPos);
	
	//timer_stop(start, stop);
	//if(newStateCountDown == 0){
		//decision_point(DIM/2, 0, nextPos, exitPos);
	//}
	//newStateCountDown--;
}

__global__ void init(int* nextPos, float *rands){
	int idx = threadIdx.x + blockIdx.x*blockDim.x;
	if(idx<AGENTNO)
		nextPos[idx]=(int)(rands[idx]*DIM*DIM);
}

__global__ void init2(int* nextPos, float*rands){
	int idx = threadIdx.x + blockIdx.x*blockDim.x;
	if(idx<AGENTNO)
		nextPos[idx]=idx;
}

__global__ void initexit(int *nextPos, int *exitPos, const int *exitList, const int exitAmount){
	int idx = threadIdx.x + blockIdx.x*blockDim.x;
	if(idx<AGENTNO){
		int t = nextPos[idx];
		int dist = POS_TO_EXIT_1D(t, exitList[0]);
		exitPos[idx] = exitList[0];
		for(int i = 1; i < exitAmount; i++){
			int dist2 = POS_TO_EXIT_1D(t, exitList[i]);
			if(dist > dist2){
				dist = dist2;
				exitPos[idx] = exitList[i];
			}
		}
	}
}

int *hexit = new int[100];
void init_state(GPUAnimBitmap &bitmap){
	int blocksize = 32;
	int gridsize = AGENTNO/blocksize+1;

	cudaMalloc(&(bitmap.rands), AGENTNO*sizeof(float));
	curandCreateGenerator(&(bitmap.gen), CURAND_RNG_PSEUDO_DEFAULT);
	curandSetPseudoRandomGeneratorSeed(bitmap.gen, 1234);
	curandGenerateUniform(bitmap.gen, bitmap.rands, AGENTNO);
	cudaMalloc((void**)&(bitmap.state1), AGENTNO*sizeof(int));
	cudaMalloc((void**)&(bitmap.state2), AGENTNO*sizeof(int));
	cudaMalloc((void**)&(bitmap.exitPos), AGENTNO*sizeof(int));
	cudaMalloc((void**)&(bitmap.env), DIM*DIM*sizeof(int));
	cudaMemset(bitmap.env, EMPTY, DIM*DIM*sizeof(int));
	init<<<gridsize,blocksize>>>(bitmap.state2 ,bitmap.rands);

	int *hdata = (int*)malloc(AGENTNO*sizeof(int));
	int *hdata2 = (int*)malloc(DIM*DIM*sizeof(int));
	memset(hdata2, 0, DIM*DIM*sizeof(int));
	cudaMemcpy(hdata, bitmap.state2, AGENTNO*sizeof(int), cudaMemcpyDeviceToHost);
	for(int i=0; i<AGENTNO; i++){
		if(hdata2[hdata[i]] == 0)
			hdata2[hdata[i]] = 1;
		else {
			int j = hdata[i];
			while(hdata2[j] != 0){
				j++;
				if(j>=DIM*DIM)
					j=0;
			}
			hdata2[j]=1;
			hdata[i]=j;
		}
		int j = hdata[i];
		//printf("pos:%d, pos_x:%d, pos_y:%d, dist:%d\n", hdata[i], hdata[i]%DIM, hdata[i]/DIM, TWO_TO_EXIT(hdata[i]%DIM,hdata[i]/DIM));
	}
	cudaMemcpy(bitmap.state2, hdata, AGENTNO*sizeof(int), cudaMemcpyHostToDevice);

	int *dexit;
	hexit[0] = TWO_TO_ONE(0, 0);
	if(EXITNO>1)
		hexit[1] = TWO_TO_ONE(0, DIM-1);
	if(EXITNO>2)
		hexit[2] = TWO_TO_ONE(DIM-1, 0);
	if(EXITNO>3)
		hexit[3] = TWO_TO_ONE(DIM-1, DIM-1);
	if(EXITNO>4)
		hexit[4] = TWO_TO_ONE(DIM/2, DIM/2);


	cudaMalloc((void**)&dexit, EXITNO*sizeof(int));
	cudaMemcpy(dexit, hexit, EXITNO*sizeof(int), cudaMemcpyHostToDevice);
	initexit<<<gridsize, blocksize>>>(bitmap.state2, bitmap.exitPos, dexit, EXITNO);
	/*
	cudaMemcpy(hdata, bitmap.exitPos, AGENTNO*sizeof(int), cudaMemcpyDeviceToHost);
	for(int i=0; i<AGENTNO; i++){
		printf("%3d ", hdata[i]);
		if(i%DIM==DIM-1)
			printf("\n");
	}*/
	cudaFree(dexit);
	free(hdata);
	free(hdata2);
}

void decision_point(int x, int y, int *nextPos, int *exitPos){
	printf("%d, %d -> ", x, y);
	int blocksize = 32;
	int gridsize = AGENTNO/blocksize+1;
	int celldim = CELLDIM;
	x = x/celldim;
	y = y/celldim;
	printf("%d, %d\n", x, y);

	hexit[EXITNO] = TWO_TO_ONE(x, y);
	int *dexit;
	cudaMalloc((void**)&dexit, (EXITNO+1)*sizeof(int));
	EXITNO++;
	cudaMemcpy(dexit, hexit, EXITNO*sizeof(int), cudaMemcpyHostToDevice);
	initexit<<<gridsize, blocksize>>>(nextPos, exitPos, dexit, EXITNO);
}
